package main

import (
	"crypto/md5"
	"crypto/sha1"
	"database/sql"
	"database/sql/driver"
	"encoding/base32"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// 默认监听IP
const IP = "0.0.0.0"

// 默认监听端口
const PORT = 6913

// sqlite数据库文件名
const DBFILE = "shortlink.db"

// 链接序号的加密密码
// 修改此密码【必须】重建数据库，而重建数据库会导致使用中的旧链接失效
const PASSWORD = "0553db9fb960e0ae452de31d52d81a9eb90a4ad8"

var dir string

var errorMsg map[string]string

var db *sql.DB

func init() {
	errorMsg = make(map[string]string, 100)
	errorMsg["ERROR_100001"] = "app_id格式错误"
	errorMsg["ERROR_100002"] = "timestamp格式错误"
	errorMsg["ERROR_100003"] = "signature格式错误"
	errorMsg["ERROR_100004"] = "应用不存在"
	errorMsg["ERROR_100005"] = "签名验证失败"
	errorMsg["ERROR_100006"] = "签名过期"
	errorMsg["ERROR_100007"] = "url格式错误"
	errorMsg["ERROR_100008"] = "expires格式错误"
	errorMsg["ERROR_100009"] = "type格式错误"
	errorMsg["ERROR_100010"] = "url不是有效的链接"
	errorMsg["ERROR_100011"] = "链接保存失败"
	errorMsg["ERROR_100012"] = "link_id格式错误"
	errorMsg["ERROR_100013"] = "链接删除失败"
}

func queryRow(object interface{}, stmt string, parameters ...interface{}) (result map[string]interface{}, err error) {
	var rows *sql.Rows
	if db, ok := object.(*sql.DB); ok {
		rows, err = db.Query(stmt, parameters...)
	} else if tx, ok := object.(*sql.Tx); ok {
		rows, err = tx.Query(stmt, parameters...)
	} else {
		err = errors.New("参数object的类型错误！")
		fmt.Println(err)
		return
	}
	if err != nil {
		fmt.Println("查询数据库失败：", err)
		fmt.Println(stmt, parameters)
		return
	}
	defer rows.Close()

	columns, _ := rows.Columns()
	valueArray := make([]interface{}, len(columns))
	pointArray := make([]interface{}, len(columns))

	for i := range valueArray {
		pointArray[i] = &valueArray[i]
	}

	if !rows.Next() {
		err = errors.New("数据表为空！")
		fmt.Println(err)
		fmt.Println(stmt, parameters)
		return
	}

	err = rows.Scan(pointArray...)
	if err != nil {
		fmt.Println("扫描数据行失败：", err)
		return
	}

	result = make(map[string]interface{})

	for i, column_name := range columns {
		result[column_name] = valueArray[i]
	}
	return
}

func execSTMT(object interface{}, stmt string, parameters ...interface{}) (result driver.Result, err error) {
	if db, ok := object.(*sql.DB); ok {
		result, err = db.Exec(stmt, parameters...)
	} else if tx, ok := object.(*sql.Tx); ok {
		result, err = tx.Exec(stmt, parameters...)
	} else {
		err = errors.New("参数object的类型错误！")
		fmt.Println(err)
		return
	}
	if err != nil {
		fmt.Println("语句执行失败：", err)
		fmt.Println(stmt, parameters)
		return
	}
	return
}

func randSecret() string {
	rand.Seed(time.Now().UnixNano())
	randomNum := rand.Uint32()
	bytes := make([]byte, 4)
	binary.BigEndian.PutUint32(bytes, randomNum)
	hash := sha1.Sum([]byte(bytes))
	hashStr := hex.EncodeToString(hash[:])
	return hashStr
}

func initSqlite() (err error) {
	_, err = os.OpenFile(dir+DBFILE, os.O_CREATE, 0755)
	if err != nil {
		fmt.Println("数据库文件打开/创建失败：", err)
		return
	}

	db, err = sql.Open("sqlite3", dir+DBFILE)
	if err != nil {
		fmt.Println("数据库连接参数格式错误：", err)
		return
	}
	defer db.Close()

	var stmt string

	stmt = `CREATE TABLE IF NOT EXISTS app (id INTEGER PRIMARY KEY AUTOINCREMENT, secret TEXT (40) DEFAULT (''), enabled INTEGER DEFAULT (1));`
	_, err = execSTMT(db, stmt)
	if err != nil {
		return
	}

	stmt = `CREATE TABLE IF NOT EXISTS link (key TEXT (20) DEFAULT ('') PRIMARY KEY, url TEXT (1000) DEFAULT (''), expires INTEGER DEFAULT (0), addtime INTEGER DEFAULT (0), type INTEGER DEFAULT (0), app_id INTEGER DEFAULT (0));`
	_, err = execSTMT(db, stmt)
	if err != nil {
		return
	}

	stmt = `CREATE TABLE IF NOT EXISTS sn (id INTEGER PRIMARY KEY AUTOINCREMENT, stage INTEGER UNIQUE);`
	_, err = execSTMT(db, stmt)
	if err != nil {
		return
	}

	stmt = `INSERT OR IGNORE INTO app(id, secret, enabled) VALUES(?, ?, ?);`
	_, err = execSTMT(db, stmt, 1, randSecret(), 1)
	if err != nil {
		return
	}

	stmt = `INSERT OR IGNORE INTO app(id, secret, enabled) VALUES(?, ?, ?);`
	_, err = execSTMT(db, stmt, 2, randSecret(), 1)
	if err != nil {
		return
	}

	stmt = `INSERT OR IGNORE INTO link(key, url, expires, addtime, type, app_id) VALUES(?, ?, ?, ?, ?, ?);`
	_, err = execSTMT(db, stmt, "0", "http://www.wpfly.cn/", 0, time.Now().Unix(), 0, 1)
	if err != nil {
		return
	}

	return
}

func responseJson(w http.ResponseWriter, code int, msg string, data interface{}) {
	w.Header().Set("Content-Type", "application/json")

	type JsonResult struct {
		Code int         `json:"code"`
		Msg  string      `json:"msg"`
		Data interface{} `json:"data"`
	}

	result, _ := json.Marshal(JsonResult{
		Code: code,
		Msg:  msg,
		Data: data,
	})

	w.Write(result)
}

func responseSuccessJson(w http.ResponseWriter, data interface{}) {
	responseJson(w, 0, "success", data)
}

func responseFailureJson(w http.ResponseWriter, code int) {
	responseJson(w, code, errorMsg[fmt.Sprintf("ERROR_%d", code)], nil)
}

func encrypt(data []byte, password string) {
	key_bytes := md5.Sum([]byte(password))
	data_len := len(data)
	for i := 0; i <= 2*data_len; i++ {
		src := i % data_len
		dst := (i + 1) % data_len
		k := i & 15
		data[dst] = byte((int(data[src])+int(data[dst]))&255) ^ key_bytes[k]
	}
}

func decrypt(data []byte, password string) {
	key_bytes := md5.Sum([]byte(password))
	data_len := len(data)
	for i := 2 * data_len; i >= 0; i-- {
		src := i % data_len
		dst := (i + 1) % data_len
		k := i & 15
		data[dst] = byte((int(data[dst]^key_bytes[k]) + 256 - int(data[src])) & 255)
	}
}

func encryptID(id uint64) string {
	bytes := make([]byte, 8)
	binary.BigEndian.PutUint64(bytes, id)
	encrypt(bytes[3:], PASSWORD)
	return base32.StdEncoding.EncodeToString(bytes[3:])
}

func decryptID(id string) (result uint64, err error) {
	bytes, err := base32.StdEncoding.DecodeString(id)
	if err != nil {
		fmt.Println("解码失败：", err)
		return
	}
	decrypt(bytes, PASSWORD)
	bytes = append(make([]byte, 3), bytes...)
	result = binary.BigEndian.Uint64(bytes)
	return
}

func verifySignature(data url.Values, signField string, secret string) bool {
	signature, ok := data[signField]
	if !ok {
		return false
	}
	delete(data, signField)
	kv := make([]string, 0, len(data))
	for k := range data {
		kv = append(kv, k)
	}
	sort.Strings(kv)
	for i, k := range kv {
		kv[i] += "=" + data[k][0]
	}
	kv = append(kv, "secret="+secret)
	hash := sha1.Sum([]byte(strings.Join(kv, "&")))
	hashStr := hex.EncodeToString(hash[:])
	fmt.Println(hashStr)
	return hashStr == strings.ToLower(signature[0])
}

func saveLink(p map[string]interface{}) (linkId string, err error) {
	tx, err := db.Begin()
	if err != nil {
		fmt.Println("事务开始失败：", err)
		return
	}

	stmt := `REPLACE INTO sn(stage) VALUES(?);`
	result, err := execSTMT(tx, stmt, 0)
	if err != nil {
		tx.Rollback()
		return
	}

	lastInsertId, _ := result.LastInsertId()
	linkId = encryptID(uint64(lastInsertId))

	stmt = `INSERT INTO link(key, url, expires, addtime, type, app_id) VALUES(?,?,?,?,?,?);`
	result, err = execSTMT(tx, stmt, linkId, p["pUrl"], p["pExpires"], time.Now().Unix(), p["pType"], p["pAppId"])
	if err != nil {
		tx.Rollback()
		return
	}

	err = tx.Commit()
	if err != nil {
		fmt.Println("事务提交失败：", err)
		tx.Rollback()
		return
	}
	return
}

func deleteLink(p map[string]interface{}) (err error) {
	stmt := `UPDATE link SET expires=-1 WHERE key=? and app_id=?;`
	result, err := execSTMT(db, stmt, p["pLinkId"], p["pAppId"])
	if err != nil {
		return
	}
	n, _ := result.RowsAffected()
	if n != 1 {
		err = errors.New("链接失效操作失败！")
		fmt.Println(err)
		return
	}
	return
}

func redirectHttpHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		re := regexp.MustCompile(`^/([A-Za-z0-9]{1,20})$`)
		matched := re.FindStringSubmatch(fmt.Sprintf("%s", r.URL))
		if matched == nil {
			w.Write([]byte("链接格式错误！"))
			return
		}

		result, err := queryRow(db, "SELECT url, expires FROM link WHERE key=?;", strings.ToUpper(matched[1]))
		if err != nil {
			w.Write([]byte("链接不存在！"))
			return
		}

		expires := result["expires"].(int64)
		if expires != 0 && expires < time.Now().Unix() {
			w.Write([]byte("链接已失效！"))
			return
		}

		http.Redirect(w, r, result["url"].(string), 302)
		return
	}
	http.NotFound(w, r)
}

func iconHttpHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, dir+"favicon.ico")
}

func linkAddHttpHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		// 获取参数
		pUrl := r.PostFormValue("url")
		pExpires := r.PostFormValue("expires")
		pType := r.PostFormValue("type")
		pAppId := r.PostFormValue("app_id")

		// 参数格式检查
		var (
			matched bool
			err     error
		)
		matched, err = regexp.MatchString(`^\S{1,1000}$`, pUrl)
		if err != nil || !matched {
			responseFailureJson(w, 100007)
			return
		}
		matched, err = regexp.MatchString(`^\d{1,10}$`, pExpires)
		if err != nil || !matched {
			responseFailureJson(w, 100008)
			return
		}
		matched, err = regexp.MatchString(`^\d$`, pType)
		if err != nil || !matched {
			responseFailureJson(w, 100009)
			return
		}
		if !strings.HasPrefix(pUrl, `http://`) && !strings.HasPrefix(pUrl, `https://`) {
			responseFailureJson(w, 100010)
			return
		}

		// 保存链接
		curlData := map[string]interface{}{
			"pUrl":     pUrl,
			"pExpires": pExpires,
			"pType":    pType,
			"pAppId":   pAppId,
		}
		link_id, err := saveLink(curlData)
		if err != nil {
			responseFailureJson(w, 100011)
			return
		}

		// 返回结果
		responseSuccessJson(w, map[string]interface{}{
			"link_id": link_id,
		})
		return
	}
	http.NotFound(w, r)
}

func linkDeleteHttpHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		// 获取参数
		pLinkId := strings.ToUpper(r.PostFormValue("link_id"))
		pAppId := r.PostFormValue("app_id")

		// 参数格式检查
		var (
			matched bool
			err     error
		)
		matched, err = regexp.MatchString(`^[2-7A-Z]{8}$`, pLinkId)
		if err != nil || !matched {
			responseFailureJson(w, 100012)
			return
		}

		// 删除链接
		curlData := map[string]interface{}{
			"pLinkId": pLinkId,
			"pAppId":  pAppId,
		}
		if deleteLink(curlData) != nil {
			responseFailureJson(w, 100013)
			return
		}

		// 返回结果
		responseSuccessJson(w, nil)
		return
	}
	http.NotFound(w, r)
}

func IDb(h http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var err error
		db, err = sql.Open("sqlite3", dir+DBFILE)
		if err != nil {
			fmt.Println("数据库连接参数格式错误：", err)
			return
		}
		defer db.Close()
		h(w, r)
	}
}

func IApi(h http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {

		r.ParseForm()
		fmt.Println(r.PostForm)

		// 获取参数
		pAppId := r.PostFormValue("app_id")
		pTimestamp := r.PostFormValue("timestamp")
		pSignature := r.PostFormValue("signature")

		// 参数格式检查
		var (
			matched bool
			err     error
		)
		matched, err = regexp.MatchString(`^\d{1,10}$`, pAppId)
		if err != nil || !matched {
			responseFailureJson(w, 100001)
			return
		}
		matched, err = regexp.MatchString(`^\d{1,10}$`, pTimestamp)
		if err != nil || !matched {
			responseFailureJson(w, 100002)
			return
		}
		matched, err = regexp.MatchString(`^[0-9a-fA-F]{40}$`, pSignature)
		if err != nil || !matched {
			responseFailureJson(w, 100003)
			return
		}

		// app_id查询
		result, err := queryRow(db, "SELECT secret FROM app WHERE id=? AND enabled=1;", pAppId)
		if err != nil {
			responseFailureJson(w, 100004)
			return
		}

		// 签名验证
		if !verifySignature(r.PostForm, "signature", result["secret"].(string)) {
			responseFailureJson(w, 100005)
			return
		}

		// 签名时效验证
		pTimestampValue, _ := strconv.ParseInt(pTimestamp, 10, 64)
		if pTimestampValue+180 < time.Now().Unix() {
			responseFailureJson(w, 100006)
			return
		}
		h(w, r)
	}
}

func fork(processType string) (cmd *exec.Cmd, err error) {
	var args []string
	isNotValue := true
	args = append(args, os.Args[0])
	for _, arg := range os.Args[1:] {
		if isNotValue {
			if !strings.HasPrefix(arg, "-") || arg == "-" || arg == "--" {
				break
			}
			if !strings.HasSuffix(arg, "=") {
				isNotValue = false
			}
		} else {
			isNotValue = true
		}
		args = append(args, arg)
	}
	args = append(args, "--type="+processType)
	cmd = &exec.Cmd{
		Path: os.Args[0],
		Args: args,
		Env:  os.Environ(),
	}
	stdout, err := os.OpenFile(dir+processType+".log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		fmt.Println("打开日志文件异常：", err)
		return
	}
	cmd.Stderr = stdout
	cmd.Stdout = stdout
	err = cmd.Start()
	if err != nil {
		fmt.Println("启动子进程失败：", err)
		return
	}
	return
}

func getCurrentDirectory() (dir string, err error) {
	ex, err := os.Executable()
	if err != nil {
		fmt.Println("获取命令路径失败：", err)
		return
	}
	dir = filepath.Dir(ex) + string(os.PathSeparator)
	return
}

func checkPortUsed(port int) bool {
	l, err := net.Listen("tcp", fmt.Sprintf(":%s", strconv.Itoa(port)))
	if err != nil {
		return true
	}
	defer l.Close()
	return false
}

func main() {
	// 打印异常
	defer func() {
		err := recover()
		if err != nil {
			fmt.Println(err)
		}
	}()

	processType := flag.String("type", "", "进程类型，默认为空")
	ip := flag.String("ip", IP, fmt.Sprintf("监听地址，默认：%s", IP))
	port := flag.Int("port", PORT, fmt.Sprintf("监听端口，默认：%d", PORT))
	flag.Parse()

    pid := os.Getpid()
	fmt.Println("进程启动：pid =", pid, ", type =", *processType)
	defer fmt.Println("进程结束：", pid)

	var err error

	dir, err = getCurrentDirectory()
	if err != nil {
		return
	}

	switch *processType {
	case "":
	
	    // 检查端口占用
		if checkPortUsed(*port) {
		    fmt.Println("端口已占用：", *port)
		    return
		}

		fmt.Println("服务器地址：", *ip, *port)

		// 主进程以异步方式，启动1个子进程（二级，守护进程）
		// 由于是异步方式，主进程在启动完子进程后就结束运行了，从而使得守护进程成为后台进程
		cmd, err := fork("daemon")
		fmt.Println("启动守护子进程：", cmd.Process.Pid)
		if err != nil {
			return
		}

	case "daemon":
	
	    // 检查端口占用
		if checkPortUsed(*port) {
		    fmt.Println("端口已占用：", *port)
		    return
		}

		// 初始化数据库
		err = initSqlite()
		if err != nil {
			return
		}

		// 守护进程以阻塞方式启动1个子进程（三级，工作进程）
		// 如果工作进程异常退出，则守护进程进入下一个循环：再次启动工作进程并阻塞。
		for i := 0; i < 10000; i++ {
			cmd, err := fork("worker")
			wPID := cmd.Process.Pid
			fmt.Println("启动工作子进程：", wPID)
			if err != nil {
				continue
			}
			cmd.Wait()
			fmt.Println("工作子进程结束：", wPID)
		}

	case "worker":

		// 工作进程
		// 必须阻塞，否则守护进程会不停地启动工作进程，直至重启次数用完。
		http.HandleFunc("/", IDb(redirectHttpHandler))
		http.HandleFunc("/favicon.ico", iconHttpHandler)
		http.HandleFunc("/api/link/add", IDb(IApi(linkAddHttpHandler)))
		http.HandleFunc("/api/link/delete", IDb(IApi(linkDeleteHttpHandler)))
		err = http.ListenAndServe(fmt.Sprintf("%s:%d", *ip, *port), nil)
		if err != nil {
		    fmt.Println("HTTP服务启动失败：", err)
			return
		}
	}
}
